"""
LangChain SQL Agent 学习示例
====================================
这个示例演示如何使用 LangChain 和大型语言模型来连接和操作 SQL 数据库。
通过自然语言与数据库进行交互，无需编写复杂的 SQL 语句。

核心概念：
1. Agent: 智能代理，理解用户意图并选择合适工具执行任务
2. Toolkit: 工具集合，提供数据库操作的各种能力
3. LLM: 大语言模型，负责理解自然语言和生成响应

学习目标：
- 掌握 LangChain 连接数据库的基本方法
- 理解 SQL Agent 的工作原理
- 学会使用自然语言查询数据库
- 了解不同的 Agent 类型及其适用场景
"""
import os

# 安装所需库（在终端中执行）:
# pip install pymysql python-dotenv langchain langchain-community langchain-openai

from dotenv import load_dotenv
from langchain.agents import AgentType
from langchain_community.agent_toolkits import create_sql_agent, SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI

"""
第一步：环境变量加载
------------------------------------
从 .env 文件加载敏感配置信息，如 API keys
.env 文件内容示例：
OPENAI_API_KEY=sk-your-openai-key-here
DATABASE_URL=mysql+pymysql://user:password@host:port/database
"""
print("正在加载环境变量...")
load_dotenv()  # 自动从当前目录的 .env 文件加载环境变量

"""
第二步：数据库连接配置
------------------------------------
使用 SQLDatabase 创建数据库连接
连接字符串格式：数据库类型+驱动://用户名:密码@主机:端口/数据库名
"""
print("正在连接数据库...")
db = SQLDatabase.from_uri(
    "mysql+pymysql://root:root@127.0.0.1:3306/fxhs_jsgl"
    # 其他数据库示例：
    # PostgreSQL: "postgresql+psycopg2://user:pass@localhost:5432/dbname"
    # SQLite: "sqlite:///./test.db"
)

"""
第四步：创建数据库工具包
------------------------------------
SQLDatabaseToolkit 提供了数据库操作的核心工具：
- sql_db_list_tables: 列出所有数据表
- sql_db_schema: 获取表结构信息
- sql_db_query: 执行SQL查询语句
- sql_db_query_checker: 检查SQL语法安全性
"""
print("正在初始化数据库工具包...")
llm = ChatOpenAI(api_key=os.getenv("XUNFEI_API_KEY"), openai_api_base=os.getenv("XUNFEI_API_BASE"), model="xdeepseekv32exp", streaming=True, temperature=0.7)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

"""
Agent 类型深度解析：

1. OPENAI_FUNCTIONS (推荐用于简单查询)
   - 优势：响应速度快，单轮决策效率高
   - 原理：使用OpenAI的函数调用功能直接选择工具
   - 适用：标准查询、数据检索、简单统计

2. ZERO_SHOT_REACT_DESCRIPTION (推荐用于复杂分析)
   - 优势：思考过程透明，适合复杂推理
   - 原理：基于ReAct框架，生成"思考→行动→观察"的循环
   - 适用：多步骤分析、复杂联接查询、需要推理的任务

3. STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION (平衡选择)
   - 优势：结合结构化和ReAct的优点
   - 原理：使用结构化输出但保持推理过程
   - 适用：需要较好工具参数处理的复杂任务

4. CONVERSATIONAL_REACT_DESCRIPTION (对话场景)
   - 优势：维护对话历史，理解上下文
   - 原理：在ReAct基础上增加对话记忆
   - 适用：多轮对话、需要上下文的连续查询
"""

"""
第六步：创建 SQL Agent
------------------------------------
Agent 是核心组件，负责：
1. 理解用户的自然语言查询
2. 规划执行步骤（思考链）
3. 选择合适的数据库工具
4. 执行并返回结果
"""
print("正在创建 SQL Agent...")
agent = create_sql_agent(
    llm=llm,  # 大语言模型，负责理解和决策
    toolkit=toolkit,  # 数据库工具集合
    verbose=False,  # 控制是否显示详细执行过程（学习时可设为True）
    agent_type=AgentType.OPENAI_FUNCTIONS  # Agent决策类型
)

"""
第七步：执行自然语言查询
------------------------------------
现在可以用自然语言与数据库交互了！
Agent会自动将自然语言转换为SQL并执行。
"""
print("开始执行查询...")

# 示例1：简单数据查询
print("\n=== 示例1：查询表数据 ===")
try:
    result = agent.invoke({"input": "查询sb_sbb表中的前5条数据"})
    print("查询结果:", result)
except Exception as e:
    print("查询出错:", e)

# 示例2：表结构查看
print("\n=== 示例2：查看表结构 ===")
try:
    result = agent.invoke({"input": "显示sb_sbb表的字段结构"})
    print("表结构查询完成")
except Exception as e:
    print("表结构查询出错:", e)

# 示例3：数据统计
print("\n=== 示例3：数据统计 ===")
try:
    result = agent.invoke({"input": "统计sb_sbb表的总记录数"})
    print("统计查询完成")
except Exception as e:
    print("统计查询出错:", e)

# 示例4：条件查询
print("\n=== 示例4：条件查询 ===")
try:
    result = agent.invoke({"input": "查询sb_sbb表中状态为'正常'的记录"})
    print("条件查询完成")
except Exception as e:
    print("条件查询出错:", e)

# 示例5：探索数据库
print("\n=== 示例5：探索数据库 ===")
try:
    result = agent.invoke({"input": "这个数据库中有哪些表？"})
    print("数据库探索完成")
except Exception as e:
    print("数据库探索出错:", e)
